"""
/***************************************************************************
Bioidx
A QGIS plugin
Calculate the biodiversity index
                                                         -------------------
begin                                : 2009-08-03 
copyright                        : (C) 2009 by Mark Baas
email                                : mark.baas123@gmail.com 
 ***************************************************************************/

/***************************************************************************
 *                                                                                                                                                 *
 *     This program is free software; you can redistribute it and/or modify    *
 *     it under the terms of the GNU General Public License as published by    *
 *     the Free Software Foundation; either version 2 of the License, or         *
 *     (at your option) any later version.                                                                     *
 *                                                                                                                                                 *
 ***************************************************************************/
"""
# Import the PyQt and QGIS libraries
from PyQt4.QtCore import * 
from PyQt4.QtGui import *
from qgis.core import *
# Initialize Qt resources from file resources.py
import resources
# Import the code for the dialog
from BioidxDialog import BioidxDialog
import os, os.path, sys

import math, tempfile
from distutils import file_util

def selectAll(pr):
    pr.select(pr.attributeIndexes(), QgsRectangle(), True)

class Bioidx: 

    def __init__(self, iface):
        # Save reference to the QGIS interface
        self.iface = iface

    def initGui(self):    
        # Create action that will start plugin configuration
        self.action = QAction(QIcon(":/plugins/habitatbioidx/icon.png"), \
                "Biodiversity by habitat", self.iface.mainWindow())
        # connect the action to the run method
        QObject.connect(self.action, SIGNAL("triggered()"), self.run) 

        # Add toolbar button and menu item
        #self.iface.addToolBarIcon(self.action)
        #self.iface.addPluginToMenu("&Biodiversity index", self.action)
        menu_bar = self.iface.mainWindow().menuBar()
       
        self.ecogis_menu = None
        
        actions = menu_bar.actions()
        
        for action in actions:
            if action.text() == "EcoGis":
                self.ecogis_menu = action.menu()

        if not self.ecogis_menu: 
            self.ecogis_menu = QMenu()
            self.ecogis_menu.setTitle("EcoGis")
            menu_bar.insertMenu(actions[ len( actions ) - 1 ], self.ecogis_menu)

        self.bio_menu = None
            
        for action in self.ecogis_menu.actions():
            if action.text() == QCoreApplication.translate("EcoGis", "Biodiversity"):
                self.bio_menu = action.menu()
                
        if not self.bio_menu:
            self.bio_menu = QMenu()
            self.bio_menu.setTitle(QCoreApplication.translate("EcoGis", "Biodiversity"))
            self.ecogis_menu.addMenu(self.bio_menu)
        
        self.bio_menu.addActions([self.action])


    def unload(self):
        # Remove the plugin menu item and icon
        pass
        #self.iface.removePluginMenu("&Biodiversity index",self.action)
        #self.iface.removeToolBarIcon(self.action)

    # run method that performs all the real work
    def run(self): 
        # create and show the dialog 
        dlg = BioidxDialog(self.dlg_ok)
        # show the dialog
        dlg.show()
        result = dlg.exec_() 
        # See if OK was pressed

    def dlg_ok(self, dlg):
            if dlg == None:
                return
    
            # do something useful (delete the line containing pass and
            # substitute with your code
            
            vl_places = dlg.getLayer("places")
            vl_species = dlg.getLayer("habitats")
            hfield = dlg.getHField()
            efield = dlg.getEField()
            sourceField = dlg.getSourceField()
            distance = dlg.getDistance()
            #output = dlg.getOutputFilename()
            pb = dlg.ui.progressBar
            
            f = dlg.getHabitatDefinitions()
            if not QFile(f).exists():
                QMessageBox.warning(None, "Biodiversity by habitat", "Habitat definitions file does not exist!")
                return False
            definitions = self.parseDefinitions(f)
            
            output = self.calculateDiversity(vl_places, vl_species, hfield, efield, sourceField, pb, definitions)

            # replace new layer for old layer

            self.addShapeToCanvas(unicode(output), vl_places)
            return True

    def parseDefinitions(self, f):
        f = open(f)
        genus = {}
        for l in f.read().split("\n"):
            if len(l.split(":")) != 2: continue
            hab = l.split(":")[0].strip()
            genus[hab] = []
            spp = l.split(":")[1].strip()
            for s in spp.split(","):
                sp = s.strip().split()[0].strip()
                if sp not in genus[hab]:
                    genus[hab].append(sp)
        return genus

    # Convinience function to add a vector layer to canvas based on input shapefile path ( as string )
    def addShapeToCanvas(self, shapeFilePath, layer):
        #shapeFilePathList = shapeFilePath.split( "/" )
        #layerName = QString( shapeFilePathList[len(shapeFilePathList)-1] )
        #if layerName.endsWith( ".shp" ):
        #    layerName = unicode( layerName ).rstrip( ".shp" )
        layerName = layer.name()
        vlayer_new = QgsVectorLayer( shapeFilePath, layerName, "ogr" )
        vlayer_new.copySymbologySettings(layer)

        if vlayer_new.isValid():
            QgsMapLayerRegistry.instance().addMapLayer(vlayer_new)
            QgsMapLayerRegistry.instance().removeMapLayer(layer.getLayerID())
            return True
        else:   
            return False

    def calculateDiversity(self, vl_places, vl_species, hfield, efield, sourceField, pb, definitions):
        # get providers
        pr_species = vl_species.dataProvider()
        pr_places = vl_places.dataProvider()
        
        # places = target layer
        output = pr_places.dataSourceUri().split("|")[0]
        output.replace("\\","/")

        # Select all
        selectAll(pr_species)
        selectAll(pr_places)

        # Create feature holders
        f_species = QgsFeature()
        f_places = QgsFeature()

        # add field to set H values
        fields = pr_places.fields()
        field = None

        for n in fields:
            o = fields[n]
            if o.name() == hfield:  field = n
        if not field:
            #print "field does not exists, adding it.."
            
            field = len(fields)
            fields[field] =  QgsField(hfield, QVariant.Double)
            fields[field+1] = QgsField(efield, QVariant.Double)
            c = 0
            #species_fields = {}
            #for hab in definitions:
            #    for sp in definitions[hab]:
            #        if not species_fields.has_key(sp): 
            #            c += 1
            #            fields[field+1+c] = QgsField(sp, QVariant.Double)
            #            species_fields[sp] = c
            
        # Setup a writer to write the new attributetabel
        if os.path.exists(output):
            # remove     
            QgsVectorFileWriter.deleteShapeFile(output)  

        writer = QgsVectorFileWriter(output, pr_places.encoding(),fields, pr_places.geometryType(), pr_places.crs() )

        # iterate through the features of places
        
        pb.setRange(0, pr_places.featureCount())
        pcnt = 0
        while pr_places.nextFeature(f_places):
            # init vars
            total_coverage = 0
            species = {}
            
            # get the geometry
            geo_places = f_places.geometry()

            # Rewind to the first feature of species
            pr_species.rewind()

            while pr_species.nextFeature(f_species):
                # get the geometry
                geo_species = f_species.geometry()

                #if  geo_places.distance(geo_species) > distance: continue
                if not geo_places.intersects(geo_species): continue
                
                # calculate the coverage in the grid
                
                geo_difference = geo_places.difference(geo_species)
                m = QgsDistanceArea()
                coverage = m.measure(geo_places)-m.measure(geo_difference)

                # there maybe be unknown values, they have to be taken off the total grid size to not interfere were the fraction (p)
                
                # get the attributes of this feature
                attrMap = f_species.attributeMap()    

                # now loop through the attributes
                for (i, attr) in attrMap.iteritems():
                    fieldName = pr_species.fields()[i].name()

                    if fieldName == sourceField:       
                        name = unicode(attr.toString())   
                        if name.strip() == "" or definitions.has_key(name) == False:
                            if name.strip() != "":
                                print "code ", name, " not found in the definition table"
                            # this grid cell should be -1
                            #species["NULL"] = None
                            break           
                        
                        # get habitats species from definitions
                        n = len(definitions[name])
                        p = float(coverage)/n # calculate AVERAGE coverage per species
                        for sp in definitions[name]:
                            if not species.has_key(sp):
                                species[sp] = 0
                            species[sp] += p
                            total_coverage += p
                        break

            # get the attributes of this feature
            attrMap = f_places.attributeMap()  
             
            # Calculate the H
            H = 0
            
            for specie in species:
                #print specie,"//", species[specie], "//"
                #fields[field+1+c] = QgsField(specie, QVariant.Double)
                #c = species_fields[specie]
                #attrMap[c] = QVariant(species[specie])
                
                p = species[specie]/total_coverage
                #print p, " = ", species[specie], " / ", total_coverage
                H += p * math.log(p) # 
                

            H = H*-1
            if total_coverage == 0: E = 0
            elif math.log(total_coverage) == 0: E = 0
            else:
                E = H/math.log(total_coverage)


            # store the H
            attrMap[field] = QVariant(H)
            attrMap[field+1] = QVariant(E)
            
            f_places.setAttributeMap(attrMap)

            writer.addFeature(f_places)
            
            pcnt += 1
            pb.setValue(pcnt)
        return output
        

